﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Linq.Expressions;
using System.Reflection;
using Test.Expressions;
using System.Text.RegularExpressions;
using Test.ExpressionVisitors;

namespace Test
{
	public class ExpressionRewriter : ExpressionVisitor
	{
		Expression root;
		int predicate_nest_level = 0;
		/// <summary>
		/// Used to pass table MemberExpressions down the tree, in order to associate the tables with any accessed fields/members.
		/// </summary>				
		Stack<LambdaExpression> lambda_stack;
		/// <summary>
		/// used primarily to append predicate JOIN/CROSS APPLY Expressions after the FROM clause
		/// </summary>
		Stack<Expression> predicatesToAppend;
		Dictionary<Expression, List<MemberExpression>> member_map;		
		Dictionary<Expression, ParameterExpression> alias_map;		
		IComparer<ColumnExpression> comparer;
		Comparison<ColumnExpression> comparison;
		public Dictionary<FieldInfo, string> field_alias_map;
		public ExpressionRewriter(Expression root) {
			this.lambda_stack = new Stack<LambdaExpression>();
			this.root = root;
			this.comparison = (a, b) => {
				return a.Name.CompareTo(b.Name);
			};
			this.comparer = new NComparer<ColumnExpression>(comparison);
			this.predicatesToAppend = new Stack<Expression>();//this.appendPredicateStack.Push(DbExpressionType.Empty);						
			this.alias_map = new Dictionary<Expression, ParameterExpression>();
			this.member_map = new Dictionary<Expression, List<MemberExpression>>();
			this.field_alias_map = new Dictionary<FieldInfo, string>();
		}
		
		public SubqueryExpression Rewrite() {
			this.root = PartialEvaluator.Eval(this.root, canEvaluateLocally);
			SubqueryExpression rewritten_expression = (SubqueryExpression)this.Visit(this.root);		
			return rewritten_expression;
		}


		protected override Expression VisitMethodCall(MethodCallExpression node) {
			if (node.Method.DeclaringType == typeof(Queryable) || node.Method.DeclaringType == typeof(Enumerable)) {
				switch (node.Method.Name) {
					case "Any":
					case "All":
						return this.WritePredicate(node);
					case "Where":
						return this.WriteWhere(node);
					case "Join":
						return this.WriteJoin(node);
					case "Count":
					case "Average":
					case "Min":
					case "Max":
					case "Sum":
						return this.WriteAggregateFunction(node);
					case "SelectMany":
						return this.WriteCrossApply(node);
					case "Take":
						return this.WriteTake(node);					
					case "Distinct":						
						return this.WriteSelect(node, true);						
					case "Select":
					default:
						return this.WriteSelect(node);
				}
			}
			else {

			}
			return node;
		}

		private Expression WriteSelect(MethodCallExpression node, bool isDistinct = false) {
			SelectExpression from;
			LambdaExpression lambda;			
			if (node.Arguments.Count > 1 && node.Arguments[1].NodeType == ExpressionType.Lambda) {
				lambda = this.GetLambda(node.Arguments[1]);
				this.alias_map.Add(node.Arguments[0], lambda.Parameters[0]);
				this.lambda_stack.Push(lambda);
				from = (SelectExpression)this.Visit(node.Arguments[0]);
				this.lambda_stack.Pop();
			}
			else {
				from = (SelectExpression)this.Visit(node.Arguments[0]);
			}			

			IEnumerable<ColumnExpression> columns = this.GetColumnsFromInnerSource(from);
			
			return new SelectExpression(
				columns: columns,
				alias: this.GetCurrentAlias(node),
				appended: null,
				from: (Expression)from,
				type: from.Type,
				where: null,
				isDistinct: isDistinct);
		
		}

		private SelectExpression WriteWhere(MethodCallExpression node) {
			LambdaExpression lambda;
			
			string alias = this.GetCurrentAlias(node);
			lambda = this.GetLambda(node.Arguments[1]);

			this.alias_map.Add(node.Arguments[0], lambda.Parameters[0]);
			this.lambda_stack.Push(lambda);			
			SelectExpression from = (SelectExpression)this.Visit(node.Arguments[0]);
			this.lambda_stack.Pop();
			Expression where = this.Visit(lambda.Body);			
		
			IEnumerable<Expression> appended = ClearAppendedStack();
			IEnumerable<ColumnExpression> columns = this.GetColumnsFromInnerSource(from);
			foreach (Expression a in appended) {
				columns = columns.Union(this.GetColumnsFromInnerSource(a), this.comparison);
			}

			return new SelectExpression(
				columns: columns,
				alias: this.GetCurrentAlias(node),
				appended: appended,
				from: (Expression)from,
				type: from.Type,
				where: where);
		}


		/// <summary>
		///This isn't a direct translation from LINQ to T-SQL. In T-SQL it wouldn't be necessary to wrap both aliases in
		///the CROSS APPLY, (one each for source and selector) into a single alias, i.e. you could still reference the original
		///2 tables. But in LINQ, the result of SelectMany is a single sequence.
		///<see cref="http://msdn.microsoft.com/en-us/library/bb534336.aspx"/>
		/// </summary>
		/// <param name="node"></param>
		/// <returns></returns>
		private Expression WriteCrossApply(MethodCallExpression node) {
			SubqueryExpression first;
			SubqueryExpression second;

			//this LambdaExpression determines what the subquery needs to select
			LambdaExpression collectionSelector;
			LambdaExpression resultSelector;			
												
			if (node.Arguments.Count == 3) {
				collectionSelector = this.GetLambda(node.Arguments[1]);
				resultSelector = this.GetLambda(node.Arguments[2]);
				ParameterExpressionRewriter rewriter = new ParameterExpressionRewriter(new Dictionary<Type, string> { { collectionSelector.Parameters[0].Type, collectionSelector.Parameters[0].Name } });
				resultSelector = (LambdaExpression)rewriter.Visit(resultSelector);

				this.alias_map.Add(node.Arguments[0], collectionSelector.Parameters[0]);
				this.alias_map.Add(collectionSelector.Body, resultSelector.Parameters[1]);//2nd Parameter of resultSelector maps to collectionSelector
				
				this.lambda_stack.Push(collectionSelector);
				this.lambda_stack.Push(resultSelector);												

				first = this.Visit(node.Arguments[0]) as SubqueryExpression;
				second = this.Visit(collectionSelector.Body) as SubqueryExpression;
				this.lambda_stack.Pop();
				this.lambda_stack.Pop();				
			}
			else {
				collectionSelector = this.GetLambda(node.Arguments[1]);
				this.alias_map.Add(node.Arguments[0], collectionSelector.Parameters[0]);
				this.lambda_stack.Push(collectionSelector);
				first = (SubqueryExpression)this.Visit(node.Arguments[0]);
				second = this.Visit(collectionSelector.Body) as SubqueryExpression;
				this.lambda_stack.Pop();				
			}


			IEnumerable<ColumnExpression> columns = this.GetColumnsFromInnerSource(from: second)
				.Union(this.GetColumnsFromInnerSource(from: first)).Distinct(comparison: this.comparison);
			var crossApply = new CrossApplyExpression(second: second, returnType: node.Method.ReturnType);

			return new SelectExpression(this.GetCurrentAlias(node),
				node.Method.ReturnType,
				columns: columns,
				from: first,
				appended: new Expression[] { crossApply });
		}

		/// <summary>
		/// Non-LambdaExpression case. Needs to get its alias from the map, if a key for it exists.
		/// </summary>
		/// <param name="node"></param>
		/// <returns></returns>
		private SelectExpression WriteJoin(MethodCallExpression node) {
			LambdaExpression firstLambda = this.GetLambda(node.Arguments[2]);
			LambdaExpression secondLambda = this.GetLambda(node.Arguments[3]);
			LambdaExpression resultSelector = this.GetLambda(node.Arguments[4]);

			ParameterExpressionRewriter rewriter = new ParameterExpressionRewriter(new Dictionary<Type, string> { 
				{ firstLambda.Parameters[0].Type, firstLambda.Parameters[0].Name },
				{ secondLambda.Parameters[0].Type, secondLambda.Parameters[0].Name }
			});
			resultSelector = (LambdaExpression)rewriter.Visit(resultSelector);
			this.alias_map.Add(node.Arguments[0], firstLambda.Parameters[0]);
			this.alias_map.Add(node.Arguments[1], secondLambda.Parameters[0]);

			this.lambda_stack.Push(resultSelector);
			this.lambda_stack.Push(firstLambda);
			this.Visit(firstLambda);						
			SubqueryExpression first = (SubqueryExpression)this.Visit(node.Arguments[0]);
			this.lambda_stack.Pop();//firstLambda
			this.lambda_stack.Push(secondLambda);
			this.Visit(secondLambda);
			SubqueryExpression second = (SubqueryExpression)this.Visit(node.Arguments[1]);
			this.lambda_stack.Pop();//secondLambda
			this.lambda_stack.Pop();//resultSelector
			Expression firstKey = this.Visit(firstLambda.Body);
			Expression secondKey = this.Visit(secondLambda.Body);
			Expression resultExpr = this.Visit(resultSelector.Body);

			var @join = new JoinExpression(
				second: second,
				returnType: resultExpr.Type,
				condition: Expression.Equal(firstKey, secondKey));

			IEnumerable<ColumnExpression> columns =
				this.GetColumnsFromInnerSource(from: first)
				.Union(this.GetColumnsFromInnerSource(from: second), this.comparison);
			return new SelectExpression(alias: this.GetCurrentAlias(node),
				type: resultExpr.Type,
				columns: columns,
				from: first,
				appended: new Expression[] { @join });
		}


		
		private Expression WritePredicate(MethodCallExpression node) {
			SubqueryExpression second;
			Expression append;
			
			LambdaExpression lambda =  (LambdaExpression)this.Visit(this.GetLambda(((MethodCallExpression)node).Arguments[1]));
			this.lambda_stack.Push(lambda);
			this.alias_map.Add(node.Arguments[0], lambda.Parameters[0]);
			second = (SubqueryExpression)this.Visit(((MethodCallExpression)node).Arguments[0]);

			second = new SubqueryExpression(type: second.Type, 
				columns:second.Columns, 
				from: second.From, 
				where: this.Visit(lambda.Body),
				take: second.Take,
				skip: second.Skip,
				appended: second.Appended,
				isDistinct: second.IsDistinct);
			
			this.lambda_stack.Pop();			
			switch (node.Method.Name) {
				case "Any":					
					return new ExistsExpression(from: second);
				case "All":
				default:
					append = new CrossApplyExpression(second: second, returnType: second.Type);
					this.predicatesToAppend.Push(append);
					return Expression.MakeBinary(ExpressionType.Equal, Expression.Constant(1), Expression.Constant(1));
			}			
			
		}
	

		private Expression WriteTake(MethodCallExpression node) {
			Expression source = node.Arguments[0];
			SelectExpression from;
			Int32 take = 0;
			if (node.Arguments.Count == 2 && node.Arguments[1].Type == typeof(Int32)) {
				take = (Int32)((ConstantExpression)node.Arguments[1]).Value;
			}
			from = (SelectExpression)this.Visit(source);
			return new SelectExpression(
				alias: this.GetCurrentAlias(node),
				columns: GetColumnsFromInnerSource(from),
				from: (Expression)from,
				take: node.Arguments[1],
				type: from.Type,
				where: null);
		}

		
		private SubqueryExpression WriteAggregateFunction(MethodCallExpression node) {
			Expression source = node.Arguments[0];
			LambdaExpression lambda;
			SelectExpression from;
			IEnumerable<ColumnExpression> columns;

			if (node == this.root) {
				from = (SelectExpression)this.Visit(source);
				return new SelectExpression(
					alias: this.GetCurrentAlias(node),
					columns: GetColumnsFromInnerSource(from),
					from: (Expression)from,
					type: from.Type,
					where: null);
			}
			switch (node.Method.Name) {
				case "Count":
					Expression where = null;
					if (node.Arguments.Count == 2) {
						lambda = this.GetLambda(node.Arguments[1]);
						this.lambda_stack.Push(lambda);
						from = (SelectExpression)this.Visit(source);
						this.lambda_stack.Pop();
						where = this.Visit(lambda.Body);
					}
					else {//if m.Arguments.Count == 1
						from = (SelectExpression)this.Visit(source);
					}
					columns = this.GetColumnsFromInnerSource(from)
						.Select(c => new ColumnExpression(c.field) {
							Alias = null,							
							Name = c.Name,
							OwnerAlias = c.OwnerAlias,
							Function = node.Method
						});
					return new AggregateSubqueryExpression(nodeType: this.GetAggregateDbExpressionType(node.Method.Name),
					returnType: node.Method.ReturnType,
					columns: columns,
					method: node.Method, from: from, where: where);
				default:
					lambda = this.GetLambda(node.Arguments[1]);//this LambdaExpression determines what the subquery needs to SELECT
					this.lambda_stack.Push(lambda);
					from = (SelectExpression)this.Visit(source);
					this.lambda_stack.Pop();
					columns = this.GetColumnsFromLambdaSelector(from.Alias, lambda)
						.Select(c => new ColumnExpression(c.field) {
							Alias = null, Name = c.Name, 
							OwnerAlias = c.OwnerAlias,
							Function = node.Method
						});
					return new AggregateSubqueryExpression(
						nodeType: this.GetAggregateDbExpressionType(node.Method.Name),
						returnType: node.Method.ReturnType,
						columns: columns,
						method: node.Method, from: from);
			}

		}

		
		protected override Expression VisitBinary(BinaryExpression node) {
			BinaryExpression placeholder;

			//effectively remove from consideration any logical expressions that use a member whose type cannot be inferred.
			if ( (node.Left.NodeType == ExpressionType.MemberAccess
				 && ((MemberExpression)node.Left).Member.DeclaringType.CanInferType() == false)
					|| (node.Left.NodeType == ExpressionType.Call && ((MethodCallExpression)node.Left).Arguments[0].NodeType == ExpressionType.MemberAccess)) {					 
					//e.g. a Boolean field, so must provide a Boolean value
					 placeholder = Expression.MakeBinary(ExpressionType.Equal, Expression.Constant(1), Expression.Constant(1));
					 if (node.NodeType.IsLogicalOperator()) {
						 return Expression.MakeBinary(node.NodeType, placeholder, this.Visit(node.Right));
					 }
					 else {
						 //it's a comparison operator
						 return placeholder;
					 }
			}
			else if ( (node.Right.NodeType == ExpressionType.MemberAccess
				&& ((MemberExpression)node.Right).Member.DeclaringType.CanInferType() == false )
				|| (node.Right.NodeType == ExpressionType.Call && ((MethodCallExpression)node.Right).Arguments[0].NodeType == ExpressionType.MemberAccess)) {
					placeholder = Expression.MakeBinary(ExpressionType.Equal, Expression.Constant(1), Expression.Constant(1));
					if (node.NodeType.IsLogicalOperator()) {						
						return Expression.MakeBinary(node.NodeType, this.Visit(node.Left), placeholder);
					}
					else {
						return placeholder;
					}				
			}
			
			return base.VisitBinary(node);
		}

		protected override Expression VisitUnary(UnaryExpression node) {
			return base.VisitUnary(node);
		}
		
		protected override Expression VisitMember(MemberExpression node) {
						
			if (node.Type.IsBoolean()) {
				return Expression.MakeBinary(ExpressionType.Equal, node, Expression.Constant(true));
			}			
			return base.VisitMember(node);
		}
		protected override Expression VisitNew(NewExpression node) {
			return base.VisitNew(node);
		}
		protected override Expression VisitMemberInit(MemberInitExpression node) {
			return base.VisitMemberInit(node);
		}
		

		/// <summary>
		/// base case - where finally the source collection is reached, or some other constant value.
		/// Non-LambdaExpression case.
		/// Instead of selecting the table directly, and handling this as a different Expression, handle all 
		/// Expressions as SelectExpression.
		/// No alias for the table within the SelectExpression, although the SelectExpression itself has an alias.
		/// </summary>
		/// <param name="node"></param>
		/// <returns></returns>
		protected override Expression VisitConstant(ConstantExpression node) {
			Type elementType;
			if (this.IsQuery(node, out elementType)) {
				IQueryable q = (IQueryable)node.Value;				
				//IEnumerable<ColumnExpression> columns = this.GetColumnsRequiredByOuterLambdaExpressions(node);
				IEnumerable<ColumnExpression> columns = this.GetAllFields(node);
				return new SelectExpression(alias: this.GetCurrentAlias(node),
					type: elementType,
					columns: columns,
					from: node);
			}
			return node;
		}

		private IEnumerable<ColumnExpression> GetAllFields(ConstantExpression from) {
			Type elementType = TypeHelper.GetElementType(from.Type);
			FieldInfo[] fields = elementType.GetFields()
				.Where(f => f.FieldType.IsScalar()).ToArray();
			ColumnExpression[] columns = new ColumnExpression[fields.Length];
			for (int i = 0; i < fields.Length; i++) {
				string alias = "field" + this.field_alias_map.Count;
				if (!field_alias_map.ContainsKey(fields[i])) {
					field_alias_map.Add(fields[i], alias);
					columns[i] = new ColumnExpression(fields[i]) {
						Alias = alias,
						OwnerAlias = null,
						Name = fields[i].Name
					};
				}				
			}
			return columns;
		}

		/// <summary>
		/// only the base case (ConstantExpression) needs to search the stack for LambdaExpressions that apply to from.
		/// </summary>
		/// <param name="from"></param>
		/// <returns></returns>
		private IEnumerable<ColumnExpression> GetColumnsRequiredByOuterLambdaExpressions(ConstantExpression from) {

			Type elementType = TypeHelper.GetElementType(from.Type);
			IEnumerable<ColumnExpression> columns = new ColumnExpression[0];
			if (this.lambda_stack.Count > 0) {
				foreach (LambdaExpression lambda_selector in this.lambda_stack) {

					//don't use the default .NET union here - not only is it slower, but it won't work.
					columns = columns.Union(MemberExpressionFinder.Find(lambda_selector.Body)
						.Where(m => m.Expression.NodeType == ExpressionType.Parameter
							&& m.Type.IsScalar()
							&& m.Member.MemberType == MemberTypes.Field
							&& (((ParameterExpression)m.Expression).Type == elementType 
							|| ((ParameterExpression)m.Expression).Type.IsAssignableFrom(elementType)))
					.Select(m => new ColumnExpression(((FieldInfo)m.Member)) {
						OwnerAlias = null,
						Name = m.Member.Name
					}), this.comparison);
					
				}
			}
			return columns.Distinct(this.comparison);
		}



		/// <summary>
		/// Re-throws/projects the columns from an inner sub-query. Assumes @from is or contains a SelectExpression.
		/// </summary>
		/// <param name="from"></param>
		/// <returns></returns>
		private IEnumerable<ColumnExpression> GetColumnsFromInnerSource(Expression @from) {
			string owner_alias;
			
			@from = SubqueryExpressionFinder.Find(@from ?? Expression.Constant(1));
			if (@from == null) {
				return new ColumnExpression[] { };
			}

			switch ((DbExpressionType)@from.NodeType) {
				case DbExpressionType.Select:
					owner_alias = ((SelectExpression)@from).Alias;
					break;
				default:
					owner_alias = null;										
					break;
			}			

			var columns = ((SubqueryExpression)@from).Columns.Select(c =>
				new ColumnExpression(c.field) {
				 OwnerAlias = owner_alias,
				 Name = c.Alias ?? c.Name,
				 //Alias = c.Alias,
				 Function = c.Function
			 });
			columns = columns.Distinct(comparison: this.comparison).ToArray();
			return columns;
		}

		private IEnumerable<ColumnExpression> GetColumnsFromLambdaSelector(string owner_alias,
			LambdaExpression lambda) {

			return MemberExpressionFinder.Find(lambda.Body)
				.Where(m => m.Expression.NodeType == ExpressionType.Parameter)
				.Select(m => new ColumnExpression(((FieldInfo)m.Member)) {
					OwnerAlias = owner_alias,
					Name = m.Member.Name
				});
		}

		int counter = 0;

		private string GetCurrentAlias(Expression node) {
			if (this.alias_map.ContainsKey(node)) {
				return this.alias_map[node].Name;
			}
			return "T" + (++counter);
		}

		IEnumerable<Expression> ClearAppendedStack() {
			List<Expression> appended = new List<Expression>();
			while (predicatesToAppend.Count > 0) {
				appended.Add(predicatesToAppend.Pop());
			}			
			return appended;
		}

		static bool canEvaluateLocally(Expression ex) {
			if (ex.NodeType == ExpressionType.Call) {
				MethodCallExpression m = ((MethodCallExpression)ex);
				for (int i = 0; i < m.Arguments.Count; i++) {
					if ((typeof(IQueryable)).IsAssignableFrom(m.Arguments[i].Type)) {
						return false;
					}
				}
				bool isqueryable = !(typeof(IQueryable)).IsAssignableFrom(ex.Type);
			}
			return ex.NodeType != ExpressionType.Parameter;
		}

	}
}
